{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2 as cv\n",
    "import numpy as np\n",
    "import apriltag\n",
    "from math import *\n",
    "from uarm.wrapper import SwiftAPI\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "swift = SwiftAPI()\n",
    "swift.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 相机内参矩阵\n",
    "mtx = np.load('./data/camera_param.npz')['mtx']\n",
    "dist = np.load('./data/camera_param.npz')['dist']\n",
    "\n",
    "# 定义一个检测器（使用字典“tag36h10”）\n",
    "at_detector = apriltag.Detector(apriltag.DetectorOptions(families='tag36h11 tag36h10'))\n",
    "\n",
    "def eulerAngles2RotationMatrix(theta):\n",
    "    \"\"\"\n",
    "    欧拉角转换旋转矩阵\n",
    "    :param theta: 欧拉角角度\n",
    "    :return: 旋转矩阵\n",
    "    \"\"\"\n",
    "    R_x = np.array([[1, 0, 0],\n",
    "                    [0, cos(theta[0]), sin(theta[0])],\n",
    "                    [0, -sin(theta[0]), cos(theta[0])]\n",
    "                    ])\n",
    "    R_y = np.array([[cos(theta[1]), 0, sin(theta[1])],\n",
    "                    [0, 1, 0],\n",
    "                    [-sin(theta[1]), 0, cos(theta[1])]\n",
    "                    ])\n",
    "    R_z = np.array([[cos(theta[2]), -sin(theta[2]),0],\n",
    "                    [sin(theta[2]), cos(theta[2]),0],\n",
    "                    [0, 0, 1]\n",
    "                    ])\n",
    "    R = R_z@R_y@R_x\n",
    "    return R\n",
    "\n",
    "def rotationMatrixToEulerAngles(R):\n",
    "    \"\"\"\n",
    "    旋转矩阵转换欧拉角\n",
    "    :param R: 旋转矩阵\n",
    "    :return: 欧拉角\n",
    "    \"\"\"\n",
    "    sy = sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])\n",
    "    singular = sy < 1e-6\n",
    "    if not singular:\n",
    "        x = atan2(R[2, 1], R[2, 2])\n",
    "        y = atan2(-R[2, 0], sy)\n",
    "        z = atan2(R[1, 0], R[0, 0])\n",
    "    else:\n",
    "        x = atan2(-R[1, 2], R[1, 1])\n",
    "        y = atan2(-R[2, 0], sy)\n",
    "        z = 0\n",
    "    return np.array([x, y, z])\n",
    "\n",
    "\n",
    "def get_tag2end(point):\n",
    "    \"\"\"\n",
    "    获取二维码到手臂末端\n",
    "    :param point:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "\n",
    "    # 计算绕z轴旋转角\n",
    "    obj= point[:,:2]\n",
    "    delta = (obj[2]-obj[3])\n",
    "    r = sqrt(delta@delta.T)\n",
    "    angle = np.arcsin(delta[0]/r)\n",
    "    # 计算变换后的坐标\n",
    "    position = swift.get_position()\n",
    "    R=eulerAngles2RotationMatrix([0,0,angle])\n",
    "    t = R @ np.array(position) - R @ point[3]\n",
    "    return R,t\n",
    "\n",
    "def get_img2tag(image,tag_size):\n",
    "    img2tag_p = []\n",
    "    cam_params = [mtx[0, 0], mtx[1, 1], mtx[0, 2], mtx[1, 2]]\n",
    "    gray = cv.cvtColor(image,cv.COLOR_BGR2GRAY)\n",
    "    tags = at_detector.detect(gray)\n",
    "    for tag in tags:\n",
    "        M, e1, e2 = at_detector.detection_pose(tag, cam_params, tag_size)\n",
    "        img2tag_p.append([tag.tag_id,tag.center,M])\n",
    "    return img2tag_p\n",
    "\n",
    "def get_base2end():\n",
    "    angle = swift.get_polar()[1]\n",
    "    position = swift.get_position()\n",
    "    R = eulerAngles2RotationMatrix([0,0,(angle-90)*pi/180])\n",
    "    t = position\n",
    "    return R,t\n",
    "\n",
    "def get_base2tag_p(tags,id,R,t):\n",
    "    ret = False\n",
    "    dst = []\n",
    "    for tag in tags:\n",
    "        if tag[0]==id:\n",
    "            pix2img_t = np.linalg.inv(mtx) @ np.insert(tag[1], 2, 1)\n",
    "            img2tag_t = (tag[2] @ np.insert(pix2img_t, 3, 1))[:3]\n",
    "            base2end_R, base2end_t = get_base2end()\n",
    "            dst = base2end_t + base2end_R @ R @ img2tag_t + base2end_R @ -R@t\n",
    "            ret = True\n",
    "            print(dst)\n",
    "    return ret,dst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "cap = cv.VideoCapture(4)\n",
    "while cap.isOpened():\n",
    "    ret,image = cap.read()\n",
    "    cv.imshow(\"image\",image)\n",
    "    if cv.waitKey(10) == 27:\n",
    "        break\n",
    "cv.destroyAllWindows()\n",
    "cap.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2 as cv\n",
    "import numpy as np\n",
    "\n",
    "def draw_circle(event,x,y,flags,param):\n",
    "    if event==cv.EVENT_LBUTTONDBLCLK:\n",
    "        cv.circle(image,(x,y),5,(255,0,0),-1) # 创建图像与窗口并将窗口与回䖲函数绑定\n",
    "\n",
    "# img=np.zeros((512,512,3),np.uint8)\n",
    "cv.namedWindow(\"image\")\n",
    "cv.setMouseCallback(\"image\",draw_circle)\n",
    "while(1):\n",
    "    cv.imshow(\"image\",image)\n",
    "    if cv.waitKey(20) == 27:\n",
    "        break\n",
    "cv.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def goto_id(image,id):\n",
    "    with np.load('./data/eyehand_Matrix.npz') as X:\n",
    "        R, t = [X[i] for i in ('R', 't')]\n",
    "    print(t)\n",
    "    tags = get_img2tag(image, 20)\n",
    "    ret, dst_p = get_base2tag_p(tags, id, R, t)\n",
    "    if ret == True:\n",
    "        print(dst_p)\n",
    "        return dst_p\n",
    "        # swift.set_position(dst_p[0], dst_p[1], int(dst_p[2]))\n",
    "    else:\n",
    "        print(\"没有识别到id%d\",id)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ -0.32265942 -67.97968943  93.05684071]\n",
      "[227.82497556  -3.77048363   8.3250682 ]\n",
      "[227.82497556  -3.77048363   8.3250682 ]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([227.82497556,  -3.77048363,   8.3250682 ])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "goto_id(image,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[277.09409057 -63.05569603   9.6964247 ]\n",
      "[277.09409057 -63.05569603   9.6964247 ]\n"
     ]
    }
   ],
   "source": [
    "dst_p = goto_id(image,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "swift.reset()\n",
    "swift.set_position(dst_p[0],dst_p[1],int(dst_p[2]+10))\n",
    "time.sleep(1)\n",
    "swift.set_position(dst_p[0],dst_p[1],int(dst_p[2]))\n",
    "time.sleep(1)\n",
    "swift.set_pump(True)\n",
    "time.sleep(1)\n",
    "swift.set_position(dst_p[0],dst_p[1],int(dst_p[2]+20))\n",
    "time.sleep(1)\n",
    "swift.set_position(dst_p[0],dst_p[1],int(dst_p[2]))\n",
    "time.sleep(1)\n",
    "swift.set_pump(False)\n",
    "time.sleep(1)\n",
    "swift.reset()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ERROR] 2024-01-09 17:01:32 [/home/lyh/PycharmProjects/vision_grab/uarm/swift/utils.py:29]: uArm is not connect\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'uArm is not connect'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "swift.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swift.set_position(100,-200,170)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'OK'"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "swift.set_servo_detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
